package me.mingshan.leetcode;

import java.util.*;

/**
 * https://leetcode.cn/problems/sum-of-distances/
 *
 * 2615. 等值距离和
 *
 * 给你一个下标从 0 开始的整数数组 nums 。现有一个长度等于 nums.length 的数组 arr 。对于满足 nums[j] == nums[i] 且 j != i 的所有 j ，
 * arr[i] 等于所有 |i - j| 之和。如果不存在这样的 j ，则令 arr[i] 等于 0 。
 *
 * 返回数组 arr 。
 *
 *
 *
 */
public class UL_2615_sum_of_distances {

    public static void main(String[] args) {

        int[] nums = {1,3,1,1,2};
        System.out.println(sumOfDistancesInArray2(nums));
    }

    public static long[] sumOfDistancesInArray(int[] nums) {
        long[] result = new long[nums.length];

        // 统计每个数字出现的索引位置，当同样的数字再出现时，不需要再次遍历了
        Map<Integer, Set<Integer>> map = new HashMap<>();

        for (int i = 0; i < nums.length; i++) {
            if (!map.containsKey(nums[i])) {
                Set<Integer> objects = new HashSet<>();
                objects.add(i);
                map.put(nums[i], objects);

                for (int j = 0; j < nums.length; j++) {
                    if (i != j && nums[i] == nums[j]) {
                        result[i] += Math.abs(i - j);

                        map.get(nums[i]).add(j);
                    }
                }
            } else {
                Set<Integer> integers = map.get(nums[i]);
                for (Integer integer : integers) {
                    result[i] += Math.abs(i - integer);
                }

            }
        }

        return result;
    }

    /**
     * 分组+前缀和
     *
     * @param nums
     * @return
     */
    public static long[] sumOfDistancesInArray2(int[] nums) {
        int n = nums.length;
        // 结果数组
        long[] ans = new long[n];
        // 存储每个值对应的下标列表
        Map<Integer, List<Integer>> groups = new HashMap<>();

        // 1. 分组存储下标
        for (int i = 0; i < n; i++) {
            groups.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
        }

        // 2. 处理每个分组
        for (List<Integer> indices : groups.values()) {
            int k = indices.size();
            // 单个元素组跳过（结果默认为0）
            if (k == 1) {
                continue;
            }

            // 3. 计算前缀和数组
            long[] prefix = new long[k + 1];
            for (int i = 0; i < k; i++) {
                prefix[i + 1] = prefix[i] + indices.get(i);
            }

            // 4. 计算每个位置的距离和
            for (int p = 0; p < k; p++) {
                int pos = indices.get(p);
                long leftSum = (long) p * pos - prefix[p];
                long rightSum = (prefix[k] - prefix[p + 1]) - (long) (k - p - 1) * pos;
                ans[pos] = leftSum + rightSum;
            }
        }

        return ans;
    }
}